import json
import requests
import base64
from datetime import datetime
import random

# 数据库账号
username = "root"
# 数据库密码
password = "root"


def get_encodeStr(username, password):
    """
    创建 URL 请求的 header
    :param username: 数据库用户名
    :param password: 用户名密码
    :return: headers
    """

    string = username + ':' + password
    encode_str = base64.b64encode(string.encode('utf-8')).decode('utf-8')
    headers = {
        'Content-Type': 'application/json',
        'Authorization': 'Basic ' + encode_str,
    }
    return headers


def get_timestamp():
    """
    创建 IotDB 接受的时间戳
    :return: timestamp
    """

    now_str = str(datetime.now().timestamp())
    timestamp = now_str.replace('.', '')
    timestamp = str(int(timestamp) + random.randint(-1000*1000*60*60*24*180,1000*1000*60*60*24*210))
    return timestamp[:-3]

def query(sql_str):
    """
    数据库查询接口
    :param sql_str: SQL 语句
    :return: 无
    """

    headers = get_encodeStr(username, password)
    data_dic = {'sql': sql_str}
    data_json = json.dumps(data_dic)
    response = requests.post('http://127.0.0.1:18080/rest/v2/query', headers=headers, data=data_json)
    if response.status_code == 200:
        print("IotDB Connection Established!")
        return response.json()
    elif response.status_code == 401:
        if response.json()['code'] == 600:
            print("IotDB Connection Error!")
            print("Wrong Login Password Error!")
            return response.json()
        elif response.json()['code'] == 603:
            print("IotDB Connection Error!")
            print("Uninitialized Auth Error!")
            return response.json()


def non_query(non_sql_str):
    """
    数据库创建、时序序列创建接口
    :param non_sql_str: non_sql 语句
    :return: 无
    """
    headers = get_encodeStr(username, password)
    data_dic = {'sql': non_sql_str}
    data_json = json.dumps(data_dic)
    response = requests.post('http://127.0.0.1:18080/rest/v2/nonQuery', headers=headers, data=data_json)
    if response.status_code == 200 and response.json()['code'] == 200:
        print("IotDB nonQuery Operation Success!")
        print(response.json())


def insert_all(iden_id, iden_sum, ser_id, data_size, src_id, dst_id, level, model, pro_type, src_ip, dst_ip, nation, route):
    """
    插入信息至 root.all_msg 数据库
    :param iden_id: 随流标识 ID
    :param iden_sum: 标识 ID 摘要
    :param ser_id: 业务数据 ID
    :param data_size: 已发送数据量
    :param src_id: 源企业 ID
    :param dst_id: 目的企业 ID
    :param level: 敏感等级
    :param model: 数据出境模式
    :param pro_type: 协议类型
    :param src_ip: 源企业 IP
    :param dst_ip: 目的企业 IP
    :param nation: 目的国家
    :param route: 路径
    :return: 插入操作信息
    """

    headers = get_encodeStr(username, password)
    timestamp = get_timestamp() # 获取当前时间的时间戳
    data_dic = {
        'timestamps': [timestamp],
        'measurements': ["iden_id", "iden_sum", "ser_id", "data_size", "src_id", "dst_id", "level", "model", "pro_type", "src_ip", "dst_ip", "nation", "route"],
        'data_types': ["TEXT", "TEXT", "TEXT", "TEXT", "TEXT", "TEXT", "TEXT", "TEXT", "TEXT", "TEXT", "TEXT", "TEXT", "TEXT"],
        'values': [[iden_id], [iden_sum], [ser_id], [data_size], [src_id], [dst_id], [level], [model], [pro_type], [src_ip], [dst_ip], [nation], [route]],
        'is_aligned': False,
        'device': "root.all_msg"
    }
    data_json = json.dumps(data_dic)
    response = requests.post('http://127.0.0.1:18080/rest/v2/insertTablet', headers=headers, data=data_json)
    if response.status_code == 200 and response.json()['code'] == 200:
        print("IotDB Insert_All Operation Success!")


def insert_company(src_id, data_size, level, model, subsidiary):
    """
    插入信息至 root.company 数据库
    :param src_id: 源企业 ID
    :param data_size: 已发送数据量
    :param level: 敏感等级
    :param model: 数据出境模式
    :return: 插入操作信息
    """

    headers = get_encodeStr(username, password)
    timestamp = get_timestamp()
    data_dic = {
        'timestamps': [timestamp],
        'measurements': ["src_id", "data_size", "level", "model", "subsidiary"],
        'data_types': ["TEXT", "TEXT", "TEXT", "TEXT", "TEXT"],
        'values': [[src_id], [data_size], [level], [model], [subsidiary]],
        'is_aligned': False,
        'device': "root.company"
    }
    data_json = json.dumps(data_dic)
    response = requests.post('http://127.0.0.1:18080/rest/v2/insertTablet', headers=headers, data=data_json)
    if response.status_code == 200 and response.json()['code'] == 200:
        print("IotDB Insert_Company Operation Success!")


def insert_province(prov_id, data_size, dst_ip):
    """
    插入信息至 root.province 数据库
    :param prov_id: 省份
    :param data_size: 已发送数据量
    :param dst_ip: 目的企业 IP
    :return:
    """
    headers = get_encodeStr(username, password)
    timestamp = get_timestamp()
    data_dic = {
        'timestamps': [timestamp],
        'measurements': ["prov_id", "data_size", "dst_ip"],
        'data_types': ["TEXT", "TEXT", "TEXT"],
        'values': [[prov_id], [data_size], [dst_ip]],
        'is_aligned': False,
        'device': "root.province"
    }
    data_json = json.dumps(data_dic)
    response = requests.post('http://127.0.0.1:18080/rest/v2/insertTablet', headers=headers, data=data_json)
    if response.status_code == 200 and response.json()['code'] == 200:
        print("IotDB Insert_Province Operation Success!")